from langchain_redis import RedisChatMessageHistory
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI

from config.load_key import load_key

history = RedisChatMessageHistory(session_id="test", redis_url=load_key("REDIS_URL"))
prompt_template = ChatPromptTemplate.from_messages([
    ("user", "{text}"),
])
llm = ChatOpenAI(
    model="Qwen/Qwen2.5-7B-Instruct",
    api_key=load_key("siliconflow_api_key"),
    base_url=load_key("siliconflow_base_url"),
)
output_parser = StrOutputParser()
chain = prompt_template | llm | output_parser
# 在链中自动进行历史记录存储
runnable = RunnableWithMessageHistory(
    chain,
    get_session_history=lambda: history,
)

# 第一次聊天，清除历史聊天记录
# history.clear()

print(runnable.invoke({"text": "请重复一次"}))
